#include "stdafx.h"

void FilePrivacy::Digest(unsigned char* source, size_t sourceSize, unsigned char** dest)
{
	HCRYPTPROV hProv;
	HCRYPTHASH hHash;

	BYTE *pbHash = NULL;
	DWORD dwHashLen = SIZE_OF_HASH; // The MD5 algorithm always returns 16 bytes. 
	DWORD cbContent = sourceSize; 
	
	string csDigest;

	if(!CryptAcquireContext(&hProv, 
		NULL, NULL, PROV_RSA_FULL, CRYPT_VERIFYCONTEXT | CRYPT_MACHINE_KEYSET)) 
	{
		return;
	}

	if(!CryptCreateHash(hProv, 
			CALG_MD5,	// algorithm identifier definitions see: wincrypt.h
			0, 0, &hHash)) 
	{
		return;
	}
	
	if(!CryptHashData(hHash, source, cbContent, 0))
	{
		return;
	}
	
	if((pbHash = (unsigned char*)malloc(dwHashLen)) == NULL)
	{
		return;
	}
	memset(pbHash, 0, dwHashLen);

	if(!CryptGetHashParam(hHash, HP_HASHVAL, pbHash, &dwHashLen, 0)) 
	{
		return;
	}

	*dest = pbHash;

	if(hHash)
	{
		CryptDestroyHash(hHash); 
	}
	if (hProv)
	{
		CryptReleaseContext(hProv, 0); 
	}
}

bool FilePrivacy::CheckPassword(LPCWSTR filePath, const wchar_t* passwd)
{
	FILE *file = NULL;
	unsigned char buf[16];
	_wfopen_s(&file, filePath, L"rb");

	if (!file)
	{
		return false;
	}

	size_t bufl = fread(buf, 1, 16, file);
	
	if (bufl!=16)
	{
		fclose(file);
		return false;
	}

	fclose(file);

	unsigned char* passHash = NULL;

	size_t length = sizeof(wchar_t) * (wcslen(passwd) + 1);
	Digest((unsigned char*)passwd, length, &passHash);

	if (passHash == NULL)
	{
		return false;
	}

	for(int i=0; i<16; i++)
	{
		if (buf[i] != passHash[i])
		{
			free(passHash);
			return false;
		}
	}

	free(passHash);
	return true;
}

PRIVRESPONSE FilePrivacy::RC4EncryptFile(HWND parent, LPCWSTR filePath, const wchar_t* passwd)
{
	// Vars
	FILE *file, *outFile;
	HCRYPTPROV hProv;
	HCRYPTHASH hHash;
	HCRYPTKEY hKey;
	unsigned long sz;
	unsigned char buf[4096];
	size_t bufl;
	
	FileMachine fm(parent, filePath); // Init FileMachine
	
	if (!fm.IsFile())
	{
		/*MessageBox(m_parent, msg.c_str(), L"Directory Encryption", MB_OK|MB_ICONHAND);*/
		return FAIL_ENC_DIR;
	}

	// Vars' Initializations
	DWORD pwlen = sizeof(wchar_t) * (wcslen(passwd) + 1);
	wstring currentFolder;
	fm.Folder(currentFolder);
	WCHAR szTempName[MAX_PATH];
	UINT dwTempFile = GetTempFileName(currentFolder.c_str(), // folder to create in
							  L"frgtmp",  // temp file name prefix 
							  0,            // create unique name 
							  szTempName);  // buffer for name 
	
	//wcout << L"created : " << szTempName << endl;
	if (dwTempFile == 0)
	{
		wcout << "GetTempFileName failed!" << endl;
		return FAIL_TEMP_FILE_NAME;
	}
	
	// Opening source and destination files
	if (_wfopen_s(&file, filePath, L"rb") != 0)
	{
		wcout << L"Error opening source file" << endl;
		return FAIL_OPEN_SRC_FILE;
	}
	if (_wfopen_s(&outFile, szTempName, L"wb") != 0)
	{
		wcout << L"Error opening dest file" << endl;
		return FAIL_OPEN_DEST_FILE;
	}

	
	// Creating hash and writing to dest file
	unsigned char* hash = NULL;
	size_t length = sizeof(wchar_t) * (wcslen(passwd) + 1);
	Digest((unsigned char*)passwd, length, &hash);

	if (hash == NULL)
	{
		// Closing files
		fclose(file);
		fclose(outFile);
		return FAIL_HASH_FAIL;
	}
	fwrite(hash, 1, SIZE_OF_HASH, outFile);
	free(hash);

	// Initializing Cryptography
	CryptAcquireContext(&hProv, NULL, MS_STRONG_PROV, PROV_RSA_FULL, 0);
	CryptCreateHash(hProv, CALG_MD5, 0, 0, &hHash);
	CryptHashData(hHash, (BYTE*)passwd, pwlen, 0);
	CryptDeriveKey(hProv, CALG_RC4, hHash, 0, &hKey);
	
	// Read & Write
	bufl=fread(buf, 1, 4096, file);
	while(bufl==4096)
	{
		sz=4096;
		CryptEncrypt(hKey, NULL, 0, 0, buf, &sz, 4096);
		fwrite(buf, 1, 4096, outFile);
		bufl=fread(buf, 1, 4096, file);
	}
	sz=bufl;
	CryptEncrypt(hKey, NULL, 1, 0, buf, &sz, 4096);
	fwrite(buf, 1, bufl, outFile);

	// Closing files
	fclose(file);
	fclose(outFile);

	// Shutting Down Cryptography
	CryptDestroyKey(hKey);
	CryptDestroyHash(hHash);
	CryptReleaseContext(hProv,0);

	// Deleting source file
	DeleteFileOperation dfo(parent, filePath, false, false, false);
	dfo.Execute();

	// Renaming dest file
	RenameFileOperation rfo(parent, szTempName, filePath, false, false);
	rfo.Execute();

	return SUCCESS_PRIV;
}

PRIVRESPONSE FilePrivacy::RC4DecryptFile(HWND parent, LPCWSTR filePath, const wchar_t* passwd)
{
	// Vars
	FILE *file, *outFile;
	HCRYPTPROV hProv;
	HCRYPTHASH hHash;
	HCRYPTKEY hKey;
	unsigned long sz;
	unsigned char buf[4096];
	size_t bufl;

	FileMachine fm(parent, filePath); // Init FileMachine
	
	// Vars' Initializations
	DWORD pwlen = sizeof(wchar_t) * (wcslen(passwd) + 1);
	wstring currentFolder;
	fm.Folder(currentFolder);
	WCHAR szTempName[MAX_PATH];
	UINT dwTempFile = GetTempFileName(currentFolder.c_str(), // folder to create in
							  L"Obg1Temp",  // temp file name prefix 
							  0,            // create unique name 
							  szTempName);  // buffer for name 
	
	//wcout << L"created : " << szTempName << endl;
	if (dwTempFile == 0)
	{
		wcout << "GetTempFileName failed!" << endl;
		return FAIL_TEMP_FILE_NAME;
	}

	// Opening source and destination files
	if (_wfopen_s(&file, filePath, L"rb") != 0)
	{
		wcout << L"Error opening source file" << endl;
		return FAIL_OPEN_DEST_FILE;
	}
	if (_wfopen_s(&outFile, szTempName, L"wb") != 0)
	{
		wcout << L"Error opening dest file" << endl;
		return FAIL_OPEN_SRC_FILE;
	}

	bufl=fread(buf, 1, 16, file);
	if (bufl!=16)
	{
		// Closing files
		fclose(file);
		fclose(outFile);
		return FAIL_READ_PASS_ERROR;
	}

	// Initializing Cryptography
	CryptAcquireContext(&hProv, NULL, MS_STRONG_PROV, PROV_RSA_FULL, 0);
	CryptCreateHash(hProv, CALG_MD5, 0, 0, &hHash);
	CryptHashData(hHash, (BYTE*)passwd, pwlen, 0);
	CryptDeriveKey(hProv, CALG_RC4, hHash, 0, &hKey);

	// Read & Write
	bufl=fread(buf, 1, 4096, file);
	while(bufl==4096)
	{
		sz=4096;
		CryptDecrypt(hKey, NULL, 0, 0, buf, &sz);
		fwrite(buf, 1, 4096, outFile);
		bufl=fread(buf, 1, 4096, file);
	}
	sz=bufl;
	CryptDecrypt(hKey, NULL, 1, 0, buf, &sz);
	fwrite(buf, 1, bufl, outFile);

	// Shutting Down Cryptography
	CryptDestroyKey(hKey);
	CryptDestroyHash(hHash);
	CryptReleaseContext(hProv,0);

	// Closing files
	fclose(file);
	fclose(outFile);

	// Deleting source file
	DeleteFileOperation dfo(parent, filePath, false, false, false);
	dfo.Execute();

	// Renaming dest file
	RenameFileOperation rfo(parent, szTempName, filePath, false, false);
	rfo.Execute();

	return SUCCESS_PRIV;
}